import re
import shutil
import sys
from copy import deepcopy
from pathlib import Path

import numpy as np
from phonopy.interface.calculator import read_crystal_structure
from phonopy.structure.atoms import PhonopyAtoms
from phonopy.structure.symmetry import Symmetry

from app.core.config import BILBAO_DATA


class MagneticAtoms(PhonopyAtoms):
    """
    一个继承自PhonopyAtoms的自定义类，用于添加磁性相关的属性以解决类型检查问题。
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.magmoms: list[float] = [0.0] * len(self)
        self.scaled_positions_conv: list[np.ndarray] = []


# 动态导入配置，支持不同的执行环境


magmoms_3d = {"Ti": 3.0, "V": 3.0, "Cr": 3.0, "Mn": 3.0, "Fe": 3.0, "Co": 3.0, "Ni": 3.0, "Cu": 3.0}
magmoms_4d = {"Y": 3.0, "Ru": 3.0, "Rh": 3.0, "Os": 2.0}
magmoms_4f = {
    "Ce": 3.0,
    "Pr": 3.0,
    "Nd": 3.0,
    "Pm": 3.0,
    "Sm": 3.0,
    "Eu": 7.0,
    "Gd": 8.0,
    "Tb": 8.0,
    "Dy": 8.0,
    "Ho": 8.0,
    "Er": 6.0,
    "Tm": 6.0,
    "Yb": 6.0,
}
magmoms = dict()
magmoms.update(magmoms_3d)
magmoms.update(magmoms_4d)
magmoms.update(magmoms_4f)
potcar = {
    "Ac": "Ac",
    "Ag": "Ag",
    "Al": "Al",
    "Ar": "Ar",
    "As": "As",
    "Au": "Au",
    "B": "B",
    "Ba": "Ba_sv",
    "Be": "Be_sv",
    "Bi": "Bi",
    "Br": "Br",
    "C": "C",
    "Ca": "Ca_sv",
    "Cd": "Cd",
    "Ce": "Ce_3",
    "Cl": "Cl",
    "Co": "Co",
    "Cr": "Cr_pv",
    "Cs": "Cs_sv",
    "Cu": "Cu_pv",
    "Dy": "Dy_3",
    "Er": "Er_3",
    "Eu": "Eu",
    "F": "F",
    "Fe": "Fe_pv",
    "Ga": "Ga_d",
    "Gd": "Gd",
    "Ge": "Ge_d",
    "H": "H",
    "He": "He",
    "Hf": "Hf_pv",
    "Hg": "Hg",
    "Ho": "Ho_3",
    "I": "I",
    "In": "In_d",
    "Ir": "Ir",
    "K": "K_sv",
    "Kr": "Kr",
    "La": "La",
    "Li": "Li_sv",
    "Lu": "Lu_3",
    "Mg": "Mg_pv",
    "Mn": "Mn_pv",
    "Mo": "Mo_pv",
    "N": "N",
    "Na": "Na_pv",
    "Nb": "Nb_pv",
    "Nd": "Nd_3",
    "Ne": "Ne",
    "Ni": "Ni_pv",
    "Np": "Np",
    "O": "O",
    "Os": "Os_pv",
    "P": "P",
    "Pa": "Pa",
    "Pb": "Pb_d",
    "Pd": "Pd",
    "Pm": "Pm_3",
    "Pr": "Pr_3",
    "Pt": "Pt",
    "Pu": "Pu",
    "Rb": "Rb_sv",
    "Re": "Re_pv",
    "Rh": "Rh_pv",
    "Ru": "Ru_pv",
    "S": "S",
    "Sb": "Sb",
    "Sc": "Sc_sv",
    "Se": "Se",
    "Si": "Si",
    "Sm": "Sm_3",
    "Sn": "Sn_d",
    "Sr": "Sr_sv",
    "Ta": "Ta_pv",
    "Tb": "Tb_3",
    "Tc": "Tc_pv",
    "Te": "Te",
    "Th": "Th",
    "Ti": "Ti_pv",
    "Tl": "Tl_d",
    "Tm": "Tm_3",
    "U": "U",
    "V": "V_pv",
    "W": "W_pv",
    "Xe": "Xe",
    "Y": "Y_sv",
    "Yb": "Yb_2",
    "Zn": "Zn",
    "Zr": "Zr_sv",
}

incar_scf = {
    "ISTART": "0",
    "ICHARG": "2",
    "ISMEAR": "-5",
    "EDIFF": "1E-6",
    "IALGO": "38",
    "PREC": "Accurate",
    "NELM": "100",
    "LORBIT": "11",
    "LWAVE": "False",
    "NPAR": "4",
    "LSORBIT": "True",
    "SAXIS": "0 0 1",
}
incar_hsp = {
    "ISTART": "1",
    "ICHARG": "11",
    "ISMEAR": "0",
    "EDIFF": "1E-6",
    "IALGO": "38",
    "PREC": "Accurate",
    "NELM": "100",
    "LORBIT": "11",
    "LWAVE": "True",
    "NPAR": "4",
    "LCHARG": "False",
    "LSORBIT": "True",
    "SAXIS": "0 0 1",
}

Pabc = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
Cabc = np.array([[0.5, 0.5, 0.0], [-0.5, 0.5, 0.0], [0.0, 0.0, 1.0]])
Cabc68 = np.array([[0.5, -0.5, 0.0], [0.5, 0.5, 0.0], [0.0, 0.0, 1.0]])
Babc = np.array([[0.5, -0.5, 0.0], [0.5, 0.5, 0.0], [0.0, 0.0, 1.0]])
Aabc = np.array([[1.0, 0.0, 0.0], [0.0, 0.5, -0.5], [0.0, 0.5, 0.5]])
Rabc = np.array([[2.0 / 3, -1.0 / 3, -1.0 / 3], [1.0 / 3, 1.0 / 3, -2.0 / 3], [1.0 / 3, 1.0 / 3, 1.0 / 3]])
Fabc = np.array([[0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0]])
Iabc = np.array([[-0.5, 0.5, 0.5], [0.5, -0.5, 0.5], [0.5, 0.5, -0.5]])


def latt_type(symbol):
    if symbol[0] == "P":
        Kc2p = Pabc
    elif symbol[0] == "C":
        Kc2p = Cabc
        if symbol == "Ccca":
            Kc2p = Cabc68
    elif symbol[0] == "B":
        Kc2p = Babc
    elif symbol[0] == "A":
        Kc2p = Aabc
    elif symbol[0] == "R":
        Kc2p = Rabc
    elif symbol[0] == "F":
        Kc2p = Fabc
    elif symbol[0] == "I":
        Kc2p = Iabc
    else:
        sys.exit("Wrong lattice type")
    return Kc2p


def v1_eq_v2(v1, v2, tol=None):
    """
    judge if vector v1 equals to vector v2 up to lattice vector
    """
    if tol is None:
        tol = 1e-4

    n = len(v1)
    eq = True
    for i in range(n):
        diff = abs(round(abs(v1[i] - v2[i])) - abs(v1[i] - v2[i]))
        if diff > tol:
            eq = False
            break
    return eq


def read_mwyck(filename: Path):
    mwyck = dict()
    with filename.open("r") as f:
        _num1, _num2, num_mwyck, num_ptrans_uni, num_ptrans_ant = list(map(int, f.readline().strip().split()))
        ptrans_uni = [np.array([0.0, 0.0, 0.0]) for ip in range(num_ptrans_uni)]
        ptrans_ant = [np.array([0.0, 0.0, 0.0]) for ip in range(num_ptrans_ant)]
        for iptrans in range(num_ptrans_uni):
            ptrans_uni[iptrans] = np.array(list(map(eval, f.readline().strip().split())))
        for iptrans in range(num_ptrans_ant):
            ptrans_ant[iptrans] = np.array(list(map(eval, f.readline().strip().split())))
        for iw in range(num_mwyck):
            tmp = f.readline().strip().split()
            mwyck[tmp[0]] = dict()
            mwyck[tmp[0]]["prim_fold"] = int(tmp[1])
            mwyck[tmp[0]]["pos_str"] = []
            mwyck[tmp[0]]["mag_str"] = []
            for istar in range(mwyck[tmp[0]]["prim_fold"]):
                line = f.readline().strip().split()
                mwyck[tmp[0]]["pos_str"].append(line[0:3])
                mwyck[tmp[0]]["mag_str"].append(line[3:])
            mwyck[tmp[0]]["pos_var"] = []
            for var in ["x", "y", "z"]:
                for i in range(3):
                    if len(re.findall(mwyck[tmp[0]]["pos_str"][0][i], var)) > 0:
                        mwyck[tmp[0]]["pos_var"].append(var)
                        break
    return mwyck, ptrans_uni, ptrans_ant


def check_mwyck(poslist, strlist, ptrans_ant, p2cR):
    rep_str = deepcopy(strlist[0])
    varlist = []
    for var in ["x", "y", "z"]:
        for i in range(3):
            if len(re.findall(rep_str[i], var)) > 0:
                varlist.append(var)
                break
    irep = 0
    match = False
    correspondence = []
    while not match and irep < len(poslist):

        correspondence = [-1 for i in range(len(poslist))]

        # assign one position to be the str representator
        rep_pos = poslist[irep]
        rep_str_num = deepcopy(strlist[0])

        changelist = dict()
        # compare position representator and str representator, get the value of variables
        if "x" in varlist:
            changelist["-2x"] = str(-2 * rep_pos[0])
            changelist["2x"] = str(2 * rep_pos[0])
            changelist["-x"] = str(-rep_pos[0])
            changelist["x"] = str(rep_pos[0])
        if "y" in varlist:
            changelist["-y"] = str(-rep_pos[1])
            changelist["y"] = str(rep_pos[1])
        if "z" in varlist:
            changelist["-z"] = str(-rep_pos[2])
            changelist["z"] = str(rep_pos[2])
        # print('changelist:')
        # print(changelist)
        for i in range(3):
            for key in changelist:
                rep_str_num[i] = rep_str_num[i].replace(key, changelist[key])
        rep_str_num[0] = eval(rep_str_num[0])
        rep_str_num[1] = eval(rep_str_num[1])
        rep_str_num[2] = eval(rep_str_num[2])

        # the position can not be the representator wyckoff, choose next one
        if not v1_eq_v2(np.dot(p2cR, rep_pos), np.dot(p2cR, rep_str_num)):
            irep += 1
            continue
        else:
            # put variable value into all strs, including anti-unitary fractional translations
            strlist_norep = [deepcopy(strlist[istr]) for istr in range(1, len(strlist))]
            strlist_num = []
            strlist_ptrans_num = []
            for i in range(len(strlist) - 1):
                for key in changelist:
                    strlist_norep[i][0] = strlist_norep[i][0].replace(key, changelist[key])
                    strlist_norep[i][1] = strlist_norep[i][1].replace(key, changelist[key])
                    strlist_norep[i][2] = strlist_norep[i][2].replace(key, changelist[key])
                strlist_norep[i][0] = eval(strlist_norep[i][0])
                strlist_norep[i][1] = eval(strlist_norep[i][1])
                strlist_norep[i][2] = eval(strlist_norep[i][2])
                strlist_num.append(np.array(strlist_norep[i]))

            # adding anti-unitary fractional translation induced mwyck
            for istr in strlist_num:
                strlist_ptrans_num.append(istr)
            if len(ptrans_ant) > 0:
                # for p in ptrans_ant:
                p = ptrans_ant[0]
                strlist_ptrans_num.append(p + np.array(rep_str_num))
                for istr in strlist_num:
                    strlist_ptrans_num.append(p + istr)

            # compare the rest str with the rest positions
            innermatch = True
            correspondence[irep] = 0
            for i in range(len(strlist_ptrans_num)):
                for j in range(len(poslist)):
                    if j != irep:
                        if v1_eq_v2(np.dot(p2cR, poslist[j]), np.dot(p2cR, strlist_ptrans_num[i])):
                            correspondence[j] = i + 1
                            break
                else:
                    innermatch = False  # cannot find one position match the str
                if not innermatch:
                    break
            else:
                match = True  # all str can find one position to match
                break  # while loop. There is no need to choose another position as representator

            if not innermatch:
                irep += 1
                continue  # while loop

    return match, correspondence


def write_magvec(num_atoms, mag_vector_onatoms, mwyck, cell_conv):
    # Note: magstrs are written under conventional basis (carefully for R lattice)
    with open("MAGVEC", "w") as w:
        with open("INCARm", "w") as ww:
            num_zeros = 0
            magmom = dict()
            for ia in range(num_atoms):
                if not mag_vector_onatoms.__contains__(ia):
                    magmom[ia] = [0.0, 0.0, 0.0]
                    w.write("{:>10.6f}  {:>10.6f}  {:>10.6f}  \n".format(0.0, 0.0, 0.0))
                else:
                    mw, order, amplitude = mag_vector_onatoms[ia]
                    num_prim_str = len(mwyck[mw]["mag_str"])
                    if order >= num_prim_str:
                        magstr = mwyck[mw]["mag_str"][order - num_prim_str]
                    else:
                        magstr = mwyck[mw]["mag_str"][order]
                    magstr_num = np.array([0.0, 0.0, 0.0])
                    for i in range(3):
                        magtmp = magstr[i]
                        magtmp = magtmp.replace("-2mx", str(-2 * amplitude))
                        magtmp = magtmp.replace("2mx", str(2 * amplitude))
                        magtmp = magtmp.replace("-2mx", str(-2 * amplitude))
                        magtmp = magtmp.replace("mx", str(amplitude))
                        magtmp = magtmp.replace("-my", str(-1 * amplitude))
                        magtmp = magtmp.replace("my", str(amplitude))
                        magtmp = magtmp.replace("-mz", str(-1 * amplitude))
                        magtmp = magtmp.replace("mz", str(amplitude))
                        magstr_num[i] = eval(magtmp)
                    magstr_num_cart = np.dot(magstr_num, cell_conv)
                    for i in range(3):
                        if order >= num_prim_str:
                            magmom[ia] = [-magstr_num_cart[0], -magstr_num_cart[1], -magstr_num_cart[2]]
                            w.write("{:>10.6f}  ".format(-magstr_num_cart[i]))
                        else:
                            magmom[ia] = [magstr_num_cart[0], magstr_num_cart[1], magstr_num_cart[2]]
                            w.write("{:>10.6f}  ".format(magstr_num_cart[i]))
                    w.write("\n")

            ww.write("MAGMOM=")
            zero_count = 0
            for ia in range(num_atoms):
                if sum(map(abs, magmom[ia])) > 1e-5:
                    if zero_count > 0:
                        ww.write("{:>3d}*0.0  ".format(zero_count))
                        zero_count = 0
                    for j in range(3):
                        if abs(magmom[ia][j]) < 1e-5:
                            ww.write(" 0 ")
                        elif abs(magmom[ia][j] - round(magmom[ia][j])) < 1e-5:
                            ww.write("{:>2d} ".format(round(magmom[ia][j])))
                        else:
                            ww.write("{:>8.5f} ".format(magmom[ia][j]))
                else:
                    zero_count += 3
            ww.write("300*0.0")
    return


def rotate_poscar(poscar: MagneticAtoms, imsg):
    og_num1 = int(msginfo[imsg]["og_ind"].split(".")[0])
    bns_num1 = int(msginfo[imsg]["bns_ind"].split(".")[0])
    sub_name = msginfo[imsg]["sub_name"]
    MA = sginfo[og_num1]["Kc2p"]
    MB = latt_type(sub_name)
    MM = sginfo[bns_num1]["Kc2p"]
    Tmat = TmatAB[imsg]
    Ttau = TtauAB[imsg]
    MIaTABMb = np.dot(np.dot(np.linalg.inv(MA), Tmat), MB)
    invT = np.linalg.inv(MIaTABMb)
    MIatau = np.dot(np.linalg.inv(MA), Ttau)
    # print(MA)
    # print(Tmat)
    # print(MB)

    # get anti-unitary fractional translations
    LCM = 96
    frac_trans = []
    enlarge = round(np.linalg.det(MIaTABMb))
    assert abs(enlarge - round(enlarge)) < 1e-5
    frac_trans_indep = [np.array([0.0, 0.0, 0.0])]
    if enlarge > 1:
        for ll in range(5):
            for mm in range(5):
                for nn in range(5):
                    inte_trans = np.array([ll + 1e-9, mm + 1e-9, nn + 1e-9])
                    frac_trans.append(np.dot(invT, inte_trans) - np.floor(np.dot(invT, inte_trans)))
        frac_trans_LCM = []
        for trans in frac_trans:
            frac_trans_LCM.append(np.round(LCM * trans).astype("int"))
        frac_trans_LCM = [tuple(i) for i in frac_trans_LCM]
        frac_trans_LCM = set(frac_trans_LCM)
        assert len(frac_trans_LCM) == enlarge
        for trans in frac_trans_LCM:
            if sum(abs(np.array(trans))) > 1e-8:
                for i in range(3):
                    if trans[i] > 0 and trans[i] < LCM:
                        frac_trans_indep.append(np.array(trans) / float(LCM))
                        break

    # define the enlarged structure
    new_cell = np.transpose(np.dot(np.transpose(poscar.cell), MIaTABMb))
    new_symbol = []
    new_positions = []
    for ind, a in enumerate(poscar.symbols):
        for trans in frac_trans_indep:
            new_symbol.append(a)
            new_scaled_position = np.dot(invT, (poscar.scaled_positions[ind] - MIatau)) + trans
            new_positions.append(new_scaled_position - np.floor(new_scaled_position))
    new_poscar = MagneticAtoms(symbols=new_symbol, scaled_positions=new_positions, cell=new_cell)

    new_poscar.scaled_positions_conv = []
    for pos in new_poscar.scaled_positions:
        new_poscar.scaled_positions_conv.append(np.dot(MM, np.array(pos)))

    return new_poscar, enlarge


def generate_incar(poscar: MagneticAtoms):
    atoms = []
    uorb = []
    uval = []
    for a in poscar.symbols:
        if a not in atoms:
            atoms.append(a)
            if a in magmoms_3d.keys() or a in magmoms_4d.keys():
                uorb.append("2")
                uval.append("3.0")
            elif a in magmoms_4f.keys():
                uorb.append("3")
                uval.append("7.0")
            else:
                uorb.append("-1")
                uval.append("0.0")
    lmaxmix = "6" if "3" in uorb else "4"
    lreal = "False" if len(poscar.symbols) <= 20 else "Auto"

    with open("INCARm", "r") as f:
        magmom = f.readline()

    with open("INCAR.scf", "w") as w:
        for key in incar_scf.keys():
            w.write(key + "=" + incar_scf[key] + "\n")
        w.write(magmom)
        w.write("\n")
        w.write("\n")
        w.write("LMAXMIX=" + lmaxmix + "\n")
        w.write("LREAL=" + lreal + "\n")
        w.write("\n")
        w.write("LDAU=.TRUE.\n")
        w.write("LDAUTYPE=2\n")
        w.write("LDAUL=")
        for l in uorb:
            w.write(l + " ")
        w.write("\n")
        w.write("LDAUU=")
        for u in uval:
            w.write(u + " ")
        w.write("\n")
        w.write("LDAUJ=")
        for u in uval:
            w.write("0.0 ")
        w.write("\n")
        w.write("LDAUPRINT=2")
    with open("INCAR.hsp", "w") as w:
        for key in incar_hsp.keys():
            w.write(key + "=" + incar_hsp[key] + "\n")
        w.write(magmom)
        w.write("\n")
        w.write("\n")
        w.write("LMAXMIX=" + lmaxmix + "\n")
        w.write("LREAL=" + lreal + "\n")
        w.write("\n")
        w.write("LDAU=.TRUE.\n")
        w.write("LDAUTYPE=2\n")
        w.write("LDAUL=")
        for l in uorb:
            w.write(l + " ")
        w.write("\n")
        w.write("LDAUU=")
        for u in uval:
            w.write(u + " ")
        w.write("\n")
        w.write("LDAUJ=")
        for u in uval:
            w.write("0.0 ")
        w.write("\n")
        w.write("LDAUPRINT=2")
    return


def generate_kpoints(subsg_ind):
    with open("KPOINTS.scf", "w") as w:
        w.write("Automatic kpoint scheme\n")
        w.write("0\n")
        w.write("auto\n")
        w.write("30\n")
    shutil.copy(str(BILBAO_DATA / "mkpoints" / f"MKPOINTS_{subsg_ind}.txt"), "KPOINTS.hsp")
    return


# ---------------------- read some global data
msginfo = dict()
with (BILBAO_DATA / "msginfo").open("r") as f:
    for i in range(1651):
        linedata = f.readline().strip().split()
        iog = int(linedata[0])
        msginfo[iog] = dict()
        msginfo[iog]["type"] = int(linedata[1])
        msginfo[iog]["bns_ind"] = linedata[2]
        msginfo[iog]["bns_name"] = linedata[3]
        msginfo[iog]["og_ind"] = linedata[4]
        msginfo[iog]["og_name"] = linedata[5]
        msginfo[iog]["sup_ind"] = int(linedata[6])
        msginfo[iog]["sup_name"] = linedata[7]
        msginfo[iog]["sub_ind"] = int(linedata[8])
        msginfo[iog]["sub_name"] = linedata[9]

sginfo = dict()
for i in range(1, 1652):
    if not sginfo.__contains__(msginfo[i]["sup_ind"]):
        sginfo[msginfo[i]["sup_ind"]] = dict()
        sginfo[msginfo[i]["sup_ind"]]["name"] = msginfo[i]["sup_name"]
for i in range(1, 231):
    sginfo[i]["Kc2p"] = latt_type(sginfo[i]["name"])

TmatAB = dict()
TtauAB = dict()
with (BILBAO_DATA / "transmat_4.txt").open("r") as f:
    for i in range(1651):
        linedata = list(map(float, f.readline().strip().split()))
        TmatAB[i + 1] = np.transpose(np.array(linedata[0:9]).reshape((3, 3)))
        TtauAB[i + 1] = np.array(linedata[9:])


def process_msg_input(origsgind: int, msgind: int):
    """
    处理磁性空间群输入，生成POSCARm、MAGVEC等文件

    Args:
        origsgind: 原始空间群编号
        msgind: 磁性空间群编号

    Returns:
        bool: 处理是否成功
    """
    # 验证参数匹配
    if origsgind != int(msginfo[msgind]["og_ind"].split(".")[0]):
        raise ValueError("Args do not match")

    bns_ind = msginfo[msgind]["bns_ind"]
    bns_num1 = int(bns_ind.split(".")[0])
    bns_num2 = int(bns_ind.split(".")[1])
    Kc2p = latt_type(msginfo[msgind]["sup_name"])
    MM = sginfo[bns_num1]["Kc2p"]
    subsg_ind = msginfo[msgind]["sub_ind"]

    # --------------------------- read mwyck
    mwyck, ptrans_uni, ptrans_ant = read_mwyck(BILBAO_DATA / "mwyck-mag" / f"{bns_num1}.{bns_num2}.txt")

    # ------------------------ original structure
    poscar_raw, _ = read_crystal_structure("POSCAR_std")
    if poscar_raw is None:
        raise ValueError("Failed to read POSCAR_std")
    poscar = MagneticAtoms(
        symbols=poscar_raw.symbols, scaled_positions=poscar_raw.scaled_positions, cell=poscar_raw.cell
    )

    poscar.scaled_positions_conv = []
    for pos in poscar.scaled_positions:
        poscar.scaled_positions_conv.append(np.dot(MM, np.array(pos)))

    new_poscar, enlarge = rotate_poscar(poscar, msgind)

    # 写入POSCARm文件
    with open("POSCARm", "w") as w:
        w.write(
            "SG#B {:>3d}   OG ({:>12s})   BNS ({:>12s})\n".format(
                subsg_ind, msginfo[msgind]["og_ind"], msginfo[msgind]["bns_ind"]
            )
        )
        w.write("{:>5.1f}\n".format(1.0))
        for i in range(3):
            w.write(
                "{:>20.14f}{:>20.14f}{:>20.14f}\n".format(
                    new_poscar.cell[i, 0], new_poscar.cell[i, 1], new_poscar.cell[i, 2]
                )
            )
        num_atoms = dict()
        for a in new_poscar.symbols:
            if not num_atoms.__contains__(a):
                num_atoms[a] = 0
            num_atoms[a] += 1
        for a in num_atoms.keys():
            w.write("{:>5s}".format(a))
        w.write("\n")
        for a in num_atoms.keys():
            w.write("{:>5d}".format(num_atoms[a]))
        w.write("\n")
        w.write("Direct\n")
        for pos in new_poscar.scaled_positions:
            w.write("{:>20.14f}{:>20.14f}{:>20.14f}\n".format(pos[0], pos[1], pos[2]))

    if enlarge > 1:
        mag_poscar = deepcopy(new_poscar)
    elif enlarge == 1:
        mag_poscar = deepcopy(poscar)
    else:
        raise ValueError(f"Unsupported enlarge factor: {enlarge}")

    mag_poscar_symm = Symmetry(mag_poscar)
    symmdata = mag_poscar_symm.dataset

    # magnetic atoms
    for ind, a in enumerate(mag_poscar.symbols):
        if a in magmoms.keys():
            mag_poscar.magmoms[ind] = magmoms[a]

    # ---------------------------- check mwycks for magnetic atoms
    mag_equivalent_atoms = []
    mag_vector_onatoms = dict()
    for ind, iatom in enumerate(symmdata.equivalent_atoms):
        if mag_poscar.magmoms[ind] > 1e-5:
            mag_equivalent_atoms.append(iatom)
    magset = set(mag_equivalent_atoms)

    for iset in magset:
        nowlist = []
        for ind, iatom in enumerate(symmdata.equivalent_atoms):
            if iatom == iset:
                nowlist.append(mag_poscar.scaled_positions_conv[ind])

        num_nowfold = len(nowlist)
        match = False
        correspondence = []
        for mw in mwyck.keys():
            antptrans = 0 if len(ptrans_ant) == 0 else 1
            if num_nowfold == mwyck[mw]["prim_fold"] * (1 + antptrans):
                for utrans in ptrans_uni:
                    nowlist_utrans = []
                    for pos in nowlist:
                        nowlist_utrans.append(pos + utrans)
                    match, correspondence = check_mwyck(
                        nowlist_utrans, mwyck[mw]["pos_str"], ptrans_ant, np.linalg.inv(Kc2p)
                    )
                    if match:
                        break

                if match:
                    icount = 0
                    for ind, iatom in enumerate(symmdata.equivalent_atoms):
                        if iatom == iset:
                            mag_vector_onatoms[ind] = (mw, correspondence[icount], mag_poscar.magmoms[ind])
                            icount += 1
                    break
        else:
            raise ValueError(f"mwyck not found for set {iset}")

    # 生成MAGVEC和KPOINTS文件
    cell_conv = np.transpose(np.dot(np.transpose(mag_poscar.cell), np.linalg.inv(MM)))
    for ind, vec in enumerate(cell_conv):
        vec_norm = np.linalg.norm(vec)
        cell_conv[ind, :] = vec / vec_norm
    write_magvec(len(mag_poscar.magmoms), mag_vector_onatoms, mwyck, cell_conv)
    generate_kpoints(subsg_ind)
